#include "gu_protocol.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <arpa/inet.h>

// Packet Management Functions

gu_packet_t* gu_packet_create(uint8_t cmd_or_resp, uint32_t payload_size) {
    gu_packet_t* packet = malloc(sizeof(gu_packet_t));
    if (!packet) return NULL;
    
    memset(packet, 0, sizeof(gu_packet_t));
    
    // Initialize header
    packet->header.length = payload_size;
    packet->header.command = cmd_or_resp;
    packet->header.status = GU_STATUS_OK;
    packet->header.flags = 0;
    
    // Allocate payload
    if (payload_size > 0) {
        packet->payload = malloc(payload_size);
        if (!packet->payload) {
            free(packet);
            return NULL;
        }
        packet->payload_size = payload_size;
        memset(packet->payload, 0, payload_size);
    }
    
    return packet;
}

void gu_packet_destroy(gu_packet_t* packet) {
    if (!packet) return;
    
    if (packet->payload) {
        free(packet->payload);
    }
    free(packet);
}

int gu_packet_send(int sockfd, gu_packet_t* packet) {
    if (!packet || sockfd < 0) return -1;
    
    // Send header
    ssize_t sent = send(sockfd, &packet->header, GU_HEADER_SIZE, 0);
    if (sent != GU_HEADER_SIZE) {
        return -1;
    }
    
    // Send payload if exists
    if (packet->header.length > 0 && packet->payload) {
        sent = send(sockfd, packet->payload, packet->header.length, 0);
        if (sent != (ssize_t)packet->header.length) {
            return -1;
        }
    }
    
    return 0;
}

gu_packet_t* gu_packet_recv(int sockfd) {
    if (sockfd < 0) return NULL;
    
    // Receive header
    gu_header_t header;
    ssize_t received = recv(sockfd, &header, GU_HEADER_SIZE, MSG_WAITALL);
    if (received != GU_HEADER_SIZE) {
        return NULL;
    }
    
    // Create packet
    gu_packet_t* packet = malloc(sizeof(gu_packet_t));
    if (!packet) return NULL;
    
    packet->header = header;
    packet->payload = NULL;
    packet->payload_size = 0;
    
    // Receive payload if exists
    if (header.length > 0) {
        if (header.length > GU_MAX_PACKET_SIZE) {
            free(packet);
            return NULL;
        }
        
        packet->payload = malloc(header.length);
        if (!packet->payload) {
            free(packet);
            return NULL;
        }
        
        received = recv(sockfd, packet->payload, header.length, MSG_WAITALL);
        if (received != (ssize_t)header.length) {
            free(packet->payload);
            free(packet);
            return NULL;
        }
        
        packet->payload_size = header.length;
    }
    
    return packet;
}

// Message Construction Functions

gu_packet_t* gu_create_connect_packet(const char* username, const char* password, const char* database) {
    if (!username || !password) return NULL;
    
    size_t username_len = strlen(username);
    size_t password_len = strlen(password);
    size_t database_len = database ? strlen(database) : 0;
    
    // Calculate payload size: 3 strings with null terminators
    uint32_t payload_size = username_len + 1 + password_len + 1 + database_len + 1;
    
    gu_packet_t* packet = gu_packet_create(GU_CMD_CONNECT, payload_size);
    if (!packet) return NULL;
    
    uint8_t* ptr = packet->payload;
    
    // Write username
    strcpy((char*)ptr, username);
    ptr += username_len + 1;
    
    // Write password
    strcpy((char*)ptr, password);
    ptr += password_len + 1;
    
    // Write database
    if (database) {
        strcpy((char*)ptr, database);
    } else {
        *ptr = '\0';
    }
    
    return packet;
}

gu_packet_t* gu_create_exec_packet(const char* sql) {
    if (!sql) return NULL;
    
    uint32_t sql_len = strlen(sql);
    gu_packet_t* packet = gu_packet_create(GU_CMD_EXEC, sql_len + 1);
    if (!packet) return NULL;
    
    strcpy((char*)packet->payload, sql);
    return packet;
}

gu_packet_t* gu_create_query_packet(const char* sql) {
    if (!sql) return NULL;
    
    uint32_t sql_len = strlen(sql);
    gu_packet_t* packet = gu_packet_create(GU_CMD_QUERY, sql_len + 1);
    if (!packet) return NULL;
    
    strcpy((char*)packet->payload, sql);
    return packet;
}

gu_packet_t* gu_create_quit_packet(void) {
    return gu_packet_create(GU_CMD_QUIT, 0);
}

gu_packet_t* gu_create_ping_packet(void) {
    return gu_packet_create(GU_CMD_PING, 0);
}

// Response Handling Functions

int gu_parse_response(gu_packet_t* packet, gu_result_t** result) {
    if (!packet) return -1;
    
    *result = NULL;
    
    switch (packet->header.command) {
        case GU_RESP_OK:
            return 0;
            
        case GU_RESP_ERROR:
            if (packet->payload) {
                fprintf(stderr, "Error: %s\n", (char*)packet->payload);
            }
            return -1;
            
        case GU_RESP_AUTH_OK:
            printf("Authentication successful\n");
            return 0;
            
        case GU_RESP_AUTH_FAIL:
            fprintf(stderr, "Authentication failed\n");
            return -1;
            
        case GU_RESP_RESULT:
            // Parse result set and create gu_result_t structure
            if (packet->payload && packet->header.length > 0) {
                *result = gu_parse_result_payload((char*)packet->payload, packet->header.length);
                if (*result) {
                    return 0;  // Success
                } else {
                    // Fallback to raw output if parsing fails
                    printf("Query result:\n%s\n", (char*)packet->payload);
                    return 0;
                }
            }
            return 0;
            
        case GU_RESP_PONG:
            printf("Pong received\n");
            return 0;
            
        default:
            fprintf(stderr, "Unknown response type: %d\n", packet->header.command);
            return -1;
    }
}

// Parse result payload from server and create gu_result_t structure
gu_result_t* gu_parse_result_payload(const char* payload, uint32_t length) {
    if (!payload || length == 0) return NULL;
    
    // printf("DEBUG CLIENT: Parsing payload length=%u, content:\n%s\n", length, payload);
    
    // Check if this is a simple status message (like "Query executed successfully.")
    if (strstr(payload, "executed successfully") || 
        strstr(payload, "OK") || 
        strstr(payload, "completed") ||
        length < 100) {  // Simple message, not table data
        
        // printf("DEBUG CLIENT: Simple message detected, returning NULL\n");
        // Just print the server message and return NULL to indicate no table data
        printf("%s", payload);
        if (payload[length-1] != '\n') {
            printf("\n");
        }
        return NULL;
    }
    
    // TODO: Parse actual table data when server returns structured results
    // For now, try to parse basic table format if it looks like table data
    
    // Look for table-like patterns (columns separated by |, rows on separate lines)
    if (strchr(payload, '|') && strchr(payload, '\n')) {
        // printf("DEBUG CLIENT: Table pattern detected, parsing...\n");
        // This might be table data, try to parse it
        // Simple implementation for now - split by lines and columns
        
        char* payload_copy = strdup(payload);
        if (!payload_copy) return NULL;
        
        char* saveptr1;
        char* line = strtok_r(payload_copy, "\n", &saveptr1);
        int line_count = 0;
        gu_result_t* result = NULL;
        
        int data_line_count = 0;
        while (line != NULL) {
            // printf("DEBUG CLIENT: Processing line %d: '%s'\n", line_count, line);
            // Skip empty lines and decorative lines (like +---+)
            if (strlen(line) > 0 && line[0] != '+') {
                // printf("DEBUG CLIENT: Line %d is data line, data_line_count=%d\n", line_count, data_line_count);
                if (data_line_count == 0 && strchr(line, '|')) {
                    // printf("DEBUG CLIENT: Processing header line\n");
                    // First data line - extract column names
                    int col_count = 0;
                    char* temp_line = strdup(line);
                    char* saveptr2;
                    char* col = strtok_r(temp_line, "|", &saveptr2);
                    while (col != NULL) {
                        col_count++;
                        col = strtok_r(NULL, "|", &saveptr2);
                    }
                    free(temp_line);
                    
                    if (col_count > 0) {
                        result = gu_result_create(col_count);
                        if (result) {
                            // Parse column names
                            temp_line = strdup(line);
                            char* saveptr3;
                            col = strtok_r(temp_line, "|", &saveptr3);
                            int col_idx = 0;
                            while (col != NULL && col_idx < col_count) {
                                // Trim whitespace
                                while (*col == ' ') col++;
                                char* end = col + strlen(col) - 1;
                                while (end > col && *end == ' ') *end-- = '\0';
                                
                                gu_result_add_column(result, col_idx, col);
                                col = strtok_r(NULL, "|", &saveptr3);
                                col_idx++;
                            }
                            free(temp_line);
                        }
                    }
                } else if (result && strchr(line, '|')) {
                    // printf("DEBUG CLIENT: Processing data row %d\n", data_line_count);
                    // Data row
                    char** row_data = malloc(result->column_count * sizeof(char*));
                    if (row_data) {
                        char* temp_line = strdup(line);
                        char* saveptr4;
                        char* col = strtok_r(temp_line, "|", &saveptr4);
                        int col_idx = 0;
                        
                        while (col != NULL && col_idx < result->column_count) {
                            // Trim whitespace
                            while (*col == ' ') col++;
                            char* end = col + strlen(col) - 1;
                            while (end > col && *end == ' ') *end-- = '\0';
                            
                            row_data[col_idx] = strdup(col);
                            col = strtok_r(NULL, "|", &saveptr4);
                            col_idx++;
                        }
                        
                        gu_result_add_row(result, row_data);
                        
                        // Free row_data array (the strings are copied by gu_result_add_row)
                        for (int i = 0; i < result->column_count; i++) {
                            free(row_data[i]);
                        }
                        free(row_data);
                        free(temp_line);
                    }
                }
                data_line_count++;
            }
            line_count++;
            line = strtok_r(NULL, "\n", &saveptr1);
            // if (line) {
            //     printf("DEBUG CLIENT: Next line found: '%s'\n", line);
            // } else {
            //     printf("DEBUG CLIENT: No more lines\n");
            // }
        }
        
        // printf("DEBUG CLIENT: Parsing completed. Result=%p\n", (void*)result);
        // if (result) {
        //     printf("DEBUG CLIENT: Result has %u columns, %u rows\n", result->column_count, result->row_count);
        // }
        
        free(payload_copy);
        return result;
    }
    
    // Fallback: just print the raw response
    printf("Server response:\n%s\n", payload);
    return NULL;
}

const char* gu_get_status_message(gu_status_t status) {
    switch (status) {
        case GU_STATUS_OK:        return "OK";
        case GU_STATUS_ERROR:     return "Error";
        case GU_STATUS_AUTH_FAIL: return "Authentication failed";
        case GU_STATUS_SYNTAX:    return "Syntax error";
        case GU_STATUS_NO_TABLE:  return "Table not found";
        case GU_STATUS_NO_DB:     return "Database not found";
        default:                  return "Unknown status";
    }
}

const char* gu_get_command_name(gu_command_t cmd) {
    switch (cmd) {
        case GU_CMD_CONNECT: return "CONNECT";
        case GU_CMD_AUTH:    return "AUTH";
        case GU_CMD_EXEC:    return "EXEC";
        case GU_CMD_QUERY:   return "QUERY";
        case GU_CMD_QUIT:    return "QUIT";
        case GU_CMD_PING:    return "PING";
        case GU_CMD_HELP:    return "HELP";
        default:             return "UNKNOWN";
    }
}

// Result Management Functions

gu_result_t* gu_result_create(uint32_t column_count) {
    if (column_count == 0) return NULL;
    
    gu_result_t* result = malloc(sizeof(gu_result_t));
    if (!result) return NULL;
    
    memset(result, 0, sizeof(gu_result_t));
    result->column_count = column_count;
    
    result->column_names = malloc(column_count * sizeof(char*));
    if (!result->column_names) {
        free(result);
        return NULL;
    }
    
    memset(result->column_names, 0, column_count * sizeof(char*));
    return result;
}

void gu_result_destroy(gu_result_t* result) {
    if (!result) return;
    
    // Free column names
    if (result->column_names) {
        for (uint32_t i = 0; i < result->column_count; i++) {
            if (result->column_names[i]) {
                free(result->column_names[i]);
            }
        }
        free(result->column_names);
    }
    
    // Free rows
    if (result->rows) {
        for (uint32_t i = 0; i < result->row_count; i++) {
            if (result->rows[i]) {
                for (uint32_t j = 0; j < result->column_count; j++) {
                    if (result->rows[i][j]) {
                        free(result->rows[i][j]);
                    }
                }
                free(result->rows[i]);
            }
        }
        free(result->rows);
    }
    
    if (result->row_lengths) {
        free(result->row_lengths);
    }
    
    free(result);
}

int gu_result_add_column(gu_result_t* result, uint32_t index, const char* name) {
    if (!result || !name || index >= result->column_count) return -1;
    
    if (result->column_names[index]) {
        free(result->column_names[index]);
    }
    
    result->column_names[index] = strdup(name);
    return result->column_names[index] ? 0 : -1;
}

int gu_result_add_row(gu_result_t* result, char** row_data) {
    if (!result || !row_data) return -1;
    
    // Reallocate rows array
    result->rows = realloc(result->rows, (result->row_count + 1) * sizeof(char**));
    if (!result->rows) return -1;
    
    result->row_lengths = realloc(result->row_lengths, (result->row_count + 1) * sizeof(uint64_t));
    if (!result->row_lengths) return -1;
    
    // Allocate new row
    result->rows[result->row_count] = malloc(result->column_count * sizeof(char*));
    if (!result->rows[result->row_count]) return -1;
    
    // Copy row data
    for (uint32_t i = 0; i < result->column_count; i++) {
        if (row_data[i]) {
            result->rows[result->row_count][i] = strdup(row_data[i]);
            result->row_lengths[result->row_count] += strlen(row_data[i]);
        } else {
            result->rows[result->row_count][i] = NULL;
        }
    }
    
    result->row_count++;
    return 0;
}

// Utility Functions

void gu_write_uint32(uint8_t* buffer, uint32_t value) {
    buffer[0] = (value >> 24) & 0xFF;
    buffer[1] = (value >> 16) & 0xFF;
    buffer[2] = (value >> 8) & 0xFF;
    buffer[3] = value & 0xFF;
}

void gu_write_uint16(uint8_t* buffer, uint16_t value) {
    buffer[0] = (value >> 8) & 0xFF;
    buffer[1] = value & 0xFF;
}

uint32_t gu_read_uint32(const uint8_t* buffer) {
    return ((uint32_t)buffer[0] << 24) |
           ((uint32_t)buffer[1] << 16) |
           ((uint32_t)buffer[2] << 8) |
           (uint32_t)buffer[3];
}

uint16_t gu_read_uint16(const uint8_t* buffer) {
    return ((uint16_t)buffer[0] << 8) | (uint16_t)buffer[1];
}

size_t gu_write_string(uint8_t* buffer, const char* str) {
    if (!str) {
        buffer[0] = '\0';
        return 1;
    }
    
    size_t len = strlen(str);
    strcpy((char*)buffer, str);
    return len + 1;
}

char* gu_read_string(const uint8_t* buffer, size_t max_len) {
    size_t len = strnlen((char*)buffer, max_len);
    if (len == max_len) len = max_len - 1;
    
    char* str = malloc(len + 1);
    if (!str) return NULL;
    
    memcpy(str, buffer, len);
    str[len] = '\0';
    return str;
}